import os
import time
import random
import copy
import numpy as np
import pandas as pd
from numpy import ndarray
from typing import List, Union
from utils import Constants, MoveRandom, Distance, LoggerConfig, plot_spectrum

logger = LoggerConfig().get_logger()
constants = Constants()
labels: List[str] = constants.labels
base_dir = constants.base_dir
data_base_dir = constants.data_base_dir
origin_path = constants.origin_path  # 原始数据存放的路径
augment_path = constants.augment_path  # 数据扩充存放的路径
each_sample_augment_count = constants.each_sample_augment_count
filenames = constants.filenames
logger.info(f"labels: {labels}")
logger.info(f"base_dir：{base_dir}")
logger.info(f"data_base_dir：{data_base_dir}")
logger.info(f"origin_path：{origin_path}")
logger.info(f"data_augment_path：{augment_path}")

img_path = constants.img_path
img_filenames = constants.img_filenames
material: str = labels[4]
origin_filename = os.path.join(origin_path, filenames["origin"][material])
augment_filename = os.path.join(augment_path, filenames["augment"][material])
origin_img_filename = os.path.join(img_path, img_filenames["origin"][material])
augment_img_filename = os.path.join(img_path, img_filenames["augment"][material])
logger.info(f"origin_filename: {origin_filename}")
logger.info(f"augment_filename: {augment_filename}")
logger.info(f"origin_img_filename: {origin_img_filename}")
logger.info(f"augment_img_filename: {augment_img_filename}")

scope_cos_distance = constants.cos_distance
scope_eucli_distance = constants.eucli_distance
logger.info(f"scope_cos_distance: {scope_cos_distance}")
logger.info(f"scope_eucli_distance: {scope_eucli_distance}")


def generator(x, nums=20):
    """
    对一个原始样本扩充为nums个样本
    :param x: 一个原始样本
    :param nums: 扩充的数据
    :return: 整体扩充的结果
    """
    start = time.perf_counter()
    count = 0
    tmp = np.empty(shape=[nums, len(x)])
    while count < nums:
        k1 = random.Random().gauss(0.0, 0.1)
        k2 = random.Random().gauss(0.0, 0.2)
        k3 = random.Random().gauss(1.0, 0.1)
        bias = np.random.normal(scale=0.03, size=x.shape) + random.Random().gauss(0.0, 0.01)
        x_ = MoveRandom(row=k1 * (x ** 3) + k2 * (x ** 2) + k3 * x + bias,
                        step=random.Random().randint(-10, 100)).random_move()
        distance = Distance(x, x_)
        cos_distance = distance.cosine()
        eucli_distance = distance.euclid()
        flag1 = scope_cos_distance["low"] < cos_distance < scope_cos_distance["high"]
        flag2 = scope_eucli_distance["low"] < eucli_distance < scope_eucli_distance["high"]
        if flag1 and flag2:
            tmp[count, :] = x_.reshape(1, -1)
            count += 1
    end = time.perf_counter()
    logger.info(f"Cost time: {(end - start):.2f} s")
    return tmp


def data_augment(origin_data: ndarray) -> ndarray:
    """
    数据扩充
    :param origin_data: 原始数据
    :return: augment_result_data: 扩充数据
    """
    logger.info("Start ...")
    start = time.perf_counter()
    _origin_data = copy.deepcopy(origin_data)
    augment_result_shape = (each_sample_augment_count * len(_origin_data), _origin_data.shape[1] - 1)
    augment_result_data: ndarray = np.empty(
        shape=augment_result_shape,
        dtype=np.float64)
    for row in range(0, len(_origin_data)):
        low = each_sample_augment_count * row
        high = each_sample_augment_count * row + each_sample_augment_count
        augment_result_data[low:high, :] = generator(
            _origin_data[row, :-1],
            nums=each_sample_augment_count)
        if each_sample_augment_count * row + each_sample_augment_count == augment_result_shape[0]:
            break
    end = time.perf_counter()
    logger.info(f" ----> Sum cost time: {(end - start):.3f} s")
    logger.info("Successfully!")

    return augment_result_data


if __name__ == '__main__':
    data = pd.read_csv(origin_filename).iloc[:, :-1].values
    plot_spectrum(data, origin_img_filename, material)

    # 数据扩充
    augment_result = data_augment(data)
    pd.DataFrame(augment_result).to_csv(augment_filename, index=False)
    plot_spectrum(augment_result[:, :], augment_img_filename, material)
    logger.info(f"Save img: {augment_filename}")
    logger.info(f"Save file: {augment_img_filename}")
